import librosa
import numpy as np
import json
import os

class FeatureExtractor:
    def __init__(self, fps=60):
        self.fps = fps
        
    def extract_features(self, audio_path):
        """
        Extract features from audio file for music visualization at specified FPS
        """
        # Check if file exists
        if not os.path.exists(audio_path):
            raise FileNotFoundError(f"Audio file '{audio_path}' not found")
            
        # Load the audio file
        y, sr = librosa.load(audio_path)
        
        # Calculate frame size based on FPS
        frame_size = int(sr / self.fps)
        
        # Extract various features
        duration = librosa.get_duration(y=y, sr=sr)
        frames = int(duration * self.fps)
        
        # Extract spectral features
        stft = np.abs(librosa.stft(y, n_fft=frame_size, hop_length=frame_size))
        
        # Spectral centroid (brightness)
        spectral_centroids = librosa.feature.spectral_centroid(y=y, sr=sr, n_fft=frame_size, hop_length=frame_size)[0]
        
        # Spectral rolloff
        spectral_rolloff = librosa.feature.spectral_rolloff(y=y, sr=sr, n_fft=frame_size, hop_length=frame_size)[0]
        
        # Spectral bandwidth
        spectral_bandwidth = librosa.feature.spectral_bandwidth(y=y, sr=sr, n_fft=frame_size, hop_length=frame_size)[0]
        
        # Zero crossing rate
        zcr = librosa.feature.zero_crossing_rate(y, frame_length=frame_size, hop_length=frame_size)[0]
        
        # RMS energy
        rms = librosa.feature.rms(y=y, frame_length=frame_size, hop_length=frame_size)[0]
        
        # Tempo and beats
        tempo, beats = librosa.beat.beat_track(y=y, sr=sr)
        
        # MFCCs (first 13)
        mfccs = librosa.feature.mfcc(y=y, sr=sr, n_mfcc=13, n_fft=frame_size, hop_length=frame_size)
        
        # Prepare features for JSON serialization
        features = {
            "duration": float(duration),
            "fps": self.fps,
            "frames": frames,
            "tempo": float(tempo),
            "features": {
                "spectral_centroids": [float(x) for x in spectral_centroids],
                "spectral_rolloff": [float(x) for x in spectral_rolloff],
                "spectral_bandwidth": [float(x) for x in spectral_bandwidth],
                "zero_crossing_rate": [float(x) for x in zcr],
                "rms": [float(x) for x in rms],
                "mfccs": [[float(x) for x in mfcc] for mfcc in mfccs]
            }
        }
        
        # Save full feature file
        base_path = os.path.splitext(audio_path)[0]
        feature_file = base_path + ".feature.json"
        
        with open(feature_file, 'w') as f:
            json.dump(features, f, indent=2)
            
        print(f"Full feature file saved to: {feature_file}")
        
        # Create example file with minimal data
        self._create_example_file(features, base_path)
        
        return features
    
    def _create_example_file(self, features, base_path):
        """
        Create a minimal example file with limited data
        """
        # Copy the structure but limit the data
        example_features = features.copy()
        example_features["features"] = {}
        
        # Take only first 100 frames for the example
        max_frames = min(100, len(features["features"]["spectral_centroids"]))
        
        for key, value in features["features"].items():
            if key == "mfccs":
                # For MFCCs, limit both frames and keep only first 5 MFCCs
                example_features["features"][key] = [
                    [float(x) for x in mfcc[:max_frames]] for mfcc in features["features"][key][:5]
                ]
            else:
                # For other features, limit to first 100 values
                example_features["features"][key] = [float(x) for x in features["features"][key][:max_frames]]
        
        example_file = base_path + ".feature.example.json"
        
        with open(example_file, 'w') as f:
            json.dump(example_features, f, indent=2)
            
        print(f"Example feature file saved to: {example_file}")